# -*- coding: utf-8 -*-
"""
BERT中文检索器
基于bert-chinese模型计算问题间的语义相似度进行检索

作者: [您的姓名]
日期: 2024
"""

import os
import numpy as np
import torch
from sklearn.metrics.pairwise import cosine_similarity
from transformers import AutoTokenizer, AutoModel
from loguru import logger
from typing import List, Dict, Tuple, Optional
from ...config import *  # 导入配置，包括镜像设置
from ...tools.cache_manager import cache_manager

# 使用全局缓存管理器
models_cache_dir = cache_manager.cache_dir


class BertChineseRetriever:
    """
    基于bert-chinese模型的相似问题检索器
    
    该检索器使用bert-chinese预训练模型来计算问题间的语义相似度，
    通过余弦相似度找到与查询问题最相似的候选问题。
    
    特点：
    - 专门针对中文优化的BERT模型
    - 使用[CLS]标记的向量表示整个句子
    - 支持自定义相似度阈值
    - 提供详细的相似度分数
    """

    def __init__(self, model_name: str = "bert-base-chinese"):
        """
        初始化BERT中文检索器

        Args:
            model_name: BERT模型名称（默认使用bert-base-chinese）
        """
        self.model_name = model_name
        self.tokenizer = None
        self.model = None
        self.sentences = []
        self.embeddings = None
        self.similarity_matrix = None
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        logger.info(f"使用设备: {self.device}")
        
        # 自动加载模型
        self.load_model()



    def load_model(self):
        """
        加载BERT模型和分词器
        """
        try:
            logger.info(f"🌐 从HuggingFace国内镜像加载模型: {self.model_name}")
            logger.info(f"🚀 镜像地址: {HF_MIRROR_ENDPOINT}")
            logger.info(f"📁 使用缓存目录: {models_cache_dir}")
            
            self.tokenizer = AutoTokenizer.from_pretrained(
                self.model_name,
                cache_dir=models_cache_dir
            )
            self.model = AutoModel.from_pretrained(
                self.model_name,
                cache_dir=models_cache_dir
            )
            self.model.eval()  # 设置为评估模式
            logger.info("✅ 模型加载成功")
        except Exception as e:
            logger.error(f"❌ 模型加载失败: {e}")
            raise

    def load_candidates(self, sentences: List[str]):
        """
        加载候选问题并计算嵌入向量

        Args:
            sentences: 候选问题列表
        """
        self.sentences = sentences
        logger.info(f"正在计算 {len(sentences)} 个候选问题的嵌入向量...")
        
        # 计算所有候选问题的嵌入向量
        self.embeddings = self.get_bert_embeddings(sentences)
        
        # 计算相似度矩阵
        self.similarity_matrix = cosine_similarity(self.embeddings)
        
        logger.info("✅ 候选问题加载完成")

    def get_bert_embeddings(self, texts: List[str]) -> np.ndarray:
        """
        获取文本的BERT嵌入向量

        Args:
            texts: 文本列表

        Returns:
            嵌入向量矩阵
        """
        embeddings = []
        
        with torch.no_grad():
            for text in texts:
                # 分词和编码
                inputs = self.tokenizer(text, return_tensors='pt', 
                                       padding=True, truncation=True, max_length=512)
                
                # 获取模型输出
                outputs = self.model(**inputs)
                
                # 使用[CLS]标记的向量作为句子表示
                cls_embedding = outputs.last_hidden_state[:, 0, :].numpy()
                embeddings.append(cls_embedding[0])
        
        return np.array(embeddings)

    def find_similar_questions(self, user_query: str, top_k: int = 5, 
                             similarity_threshold: float = 0.0) -> List[Dict]:
        """
        查找与用户查询最相似的top_k个问题

        Args:
            user_query: 用户查询的文本
            top_k: 返回的结果数量
            similarity_threshold: 相似度阈值

        Returns:
            相似问题列表
        """
        if not self.sentences:
            logger.error("❌ 候选问题未加载")
            return []

        # 获取查询的嵌入向量
        query_embedding = self.get_bert_embeddings([user_query])[0]

        # 计算查询与所有候选问题的相似度
        similarities = cosine_similarity([query_embedding], self.embeddings)[0]

        # 获取相似度最高的top_k个问题
        top_indices = np.argsort(similarities)[::-1][:top_k]
        results = []
        for idx in top_indices:
            if similarities[idx] >= similarity_threshold:
                results.append({
                    'question': self.sentences[idx],
                    'similarity': float(similarities[idx])
                })

        return results

    def batch_similarity(self, queries: List[str], top_k: int = 5) -> Dict[str, List[Dict]]:
        """
        批量计算多个查询的相似问题

        Args:
            queries: 查询问题列表
            top_k: 每个查询返回的相似问题数量

        Returns:
            字典，键为查询，值为相似问题列表
        """
        results = {}
        for query in queries:
            results[query] = self.find_similar_questions(query, top_k=top_k)
        return results

    def get_similarity_matrix(self) -> np.ndarray:
        """返回相似度矩阵"""
        return self.similarity_matrix

    def save_embeddings(self, filepath: str):
        """保存嵌入向量到文件"""
        np.save(filepath, self.embeddings)

    def load_embeddings(self, filepath: str):
        """从文件加载嵌入向量"""
        self.embeddings = np.load(filepath)


if __name__ == "__main__":
    # 从统一数据文件加载问题数据
    try:
        from questionretrieval.io import load_sample_questions, get_test_queries
        questions = load_sample_questions(5)  # 加载5个示例问题
        query = get_test_queries(questions, 1)[0]  # 生成1个测试查询
        logger.info(f"✅ 成功加载 {len(questions)} 个候选问题")
    except Exception as e:
        logger.error(f"❌ 加载问题数据失败: {e}")
        exit(1)
    
    # 创建推荐器并加载候选问题
    retriever = BertChineseRetriever()
    retriever.load_candidates(questions)
        
    # 查找相似问题
    results = retriever.find_similar_questions(query, top_k=3)

    logger.info(f"查询: {query}")
    logger.info("相似问题:")
    for result in results:
        logger.info(f"  - {result['question']} (相似度: {result['similarity']:.3f})")